Sparse linear algebraic methods and MPI for spatio-temporal inference

Bayesian Hierarchical Model

$\def\bold#1{\mathbf #1}$ $\def\yvec{\bold{y}}$ $\def\svec{\bold{s}}$ $\def\vvec{\bold{v}}$ $\def\wvec{\bold{w}}$ $\def\gammab{\bold{\gamma}}$ $\def\thetab{\bold{\theta}}$ $\def\kvec{\bold{k}}$ $\def\phib{\bold{\phi}}$ $\def\psib{\bold{\psi}}$ $\def\evec{\bold{e}}$ $\def\xvec{\bold{x}}$ $\def\Qmat{\bold{Q}}$ $\def\Amat{\bold{A}}$ $\def\betab{\bold{\beta}}$ $\def\Lmat{\bold{L}}$ $\def\Thetab{\bold{\Theta}}$ $\def\Jmat{\bold{J}}$ $\def\Cmat{\bold{C}}$ $\def\Tmat{\bold{T}}$

Divide an conquer approach to a complext problem.

$\yvec_t$ : observables
$X_t(\svec)$ : spatio-temporal process
$\vvec_t$ : observation error
$\wvec_t$ : random forcing
$\thetab_i$ : process/observation parameters
$\gammab_i$ : hyper-parameters

Bayesian Hierarchical Model

L1: Observation model:

y t = g (X t (s), v t; θ h)

$\yvec_t = g(X_t(\svec), \vvec_t; \thetab_h)$

L2: Process model:

X t (s) = f (s, t, w t (s); θ m)

$X_t(\svec) = f(\svec,t,\wvec_t(\svec); \thetab_m)$

L3: Parameter model:

p (θ m; γ m), p (θ h; γ h), p (θ v; γ v), p (θ w; γ w)

$p(\thetab_m ; \gammab_m), p(\thetab_h; \gammab_h), p(\thetab_v ; \gammab_v), p(\thetab_w; \gammab_w)$

Modelling approaches

Two approaches:

Geo-statistical models:

X t (s) \sim G P (μ t (s), K (t, t', s, s'))

$X_t(\svec) \sim GP(\mu_t(\svec),K(t,t',\svec,\svec'))$

General-purpose model, used when the dynamic structure is largely unknown.

Dynamic (mechanistically-motivated) models

X t (s) = f (X t - 1 (s), X t - 2 (s), \dots, w t (s); θ m)

$X_t(\svec) = f(X_{t-1}(\svec), X_{t-2}(\svec), \dots, \wvec_t(\svec); \thetab_m)$

Examples:
- STAR models
- Integro-difference equation models
- PDE models
- CMLs

Dynamic process models

Used when we have some mechanistic insight of the underlying process.
In the environmental sciences these are usually in the form of (stochastic) PDEs.
They implicitly define very complicated spatio-temporal covariances.

\partial X ( s ; t ) \partial t - β \partial 2 X ( s ; t ) \partial s 2 + a X (s; t) = e (s; t)

$\frac{\partial X(s;t)}{\partial t} - \beta \frac{\partial^2X(s;t)}{\partial s^2} + a X(s;t) = e(s;t)$

where $\alpha >0, \beta > 0$ and $e(s;t)$ is a random, zero mean error process. Then the geo-statistical approach requires use of covariance function:

Dynamic process models

K (u; v) = K ( 0 ; 0 ) 2 ⎛ ⎝ ⎜ ⎜ ⎜ e - u (α β) 1 2 E r f c ⎛ ⎝ ⎜ ⎜ ⎜ 2 v ( α β ) 1 2 - u β 2 ( v β ) 1 2 ⎞ ⎠ ⎟ ⎟ ⎟ + e u (α β) 1 2 E r f c ⎛ ⎝ ⎜ ⎜ ⎜ 2 v ( α β ) 1 2 + u β 2 ( v β ) 1 2 ⎞ ⎠ ⎟ ⎟ ⎟ ⎞ ⎠ ⎟ ⎟ ⎟

$K(u;v) = \frac{K(0;0)}{2}\left(e^{-u^{\left( \frac{\alpha}{\beta}\right)^\frac{1}{2}}} Erfc\left(\frac{2v\left( \frac{\alpha}{\beta}\right)^\frac{1}{2} - \frac{u}{\beta}}{2\left( \frac{v}{\beta}\right)^\frac{1}{2}} \right) + \\ e^{u^{\left( \frac{\alpha}{\beta}\right)^\frac{1}{2}}} Erfc\left(\frac{2v\left( \frac{\alpha}{\beta}\right)^\frac{1}{2} + \frac{u}{\beta}}{2\left( \frac{v}{\beta}\right)^\frac{1}{2}}\right) \right)$

(Wikle, LGM2013, Reykjavik, Iceland)

Moving on to discrete time

Continuous-time ST processes are computationally tedious to deal with.
Focus on models of the form

X t (s) = f (X t - 1 (s)) + e t (s)

$X_t(\svec) = f(X_{t-1}(\svec)) + e_t(\svec)$

where $~e_t(\svec)~$ is the solution to

(κ 2 - Δ) α / 2 (τ e (s)) = W (s)

$(\kappa^2 - \Delta)^{\alpha/2}(\tau e(\svec)) = \mathcal{W(\svec)}$

where $\mathcal{W(\svec)}$ is a space-time (white) Wiener process.

SPDE link to Matern fields

The SPDE

(κ 2 - Δ) α / 2 (τ e (s)) = W (s)

$(\kappa^2 - \Delta)^{\alpha/2}(\tau e(\svec)) = \mathcal{W(\svec)}$

has as solution (Whittle, 1954)

e (s) \sim G P (0, K (s, s'; α, τ, κ))

$e(\svec) \sim \mathcal{GP}(0,K(\svec,\svec'; \alpha,\tau,\kappa))$

where $k(\cdot,\cdot)$ is a Matérn covariance function. The wave spectrum of $e(\svec)$ is given by

R (k) = (2 π) - α (κ 2 + ∥ k ∥ 2) - α

$R(\kvec) = (2\pi)^{-\alpha}(\kappa^2 + \| \kvec \|^2)^{-\alpha}$

Why SPDEs

Non-stationarity can be handled by making the SPDE spatially varying. This results in a covariance function which is not a Matern field, but do we care?
We can model SPDEs on manifolds, not restricted to $\mathbb{R}^2$ .
SPDEs can be approximated to arbitrary precision using an appropriate basis (Lindgren, 2011).

Dimensionality reduction

X t + 1 (s) = A X t (s) + e t s

$X_{t+1}(\svec) = \mathcal{A}X_t(\svec) + e_t{\svec}$

Project:

⟨ ϕ, X t + 1 (s) ⟩ = ⟨ ϕ, A X t + 1 (s) ⟩ + ⟨ ϕ, e t (s) ⟩

$\langle \phib,X_{t+1}(\svec) \rangle = \langle \phib,\mathcal{A}X_{t+1}(\svec) \rangle + \langle \phib,e_t(\svec)\rangle$

Interpolate:

⟨ ϕ, ψ T ⟩ x t + 1 = ⟨ ϕ, A ψ T ⟩ x t + ⟨ ϕ, ψ T ⟩ e t

$\langle \phib,\psib^T \rangle\xvec_{t+1} = \langle \phib,\mathcal{A}\psib^T \rangle\xvec_t + \langle \phib,\psib^T\rangle\evec_t$

For $\alpha = 2$ and $\psib = \phib$ (Galerkin method):

[⟨ ϕ, κ 2 ϕ T ⟩ - ⟨ ϕ, Δ ϕ T ⟩)] τ e = ⟨ ϕ, W (s) ⟩

$[ \langle\phib, \kappa^2\phib^T\rangle - \langle\phib, \Delta\phib^T\rangle)]\tau \evec = \langle \phib,\mathcal{W(\svec)} \rangle$

Dimensionality reduction

e \sim N (0, Q - 1)

$\Large \evec \sim \mathcal{N}(\bold{0},\Qmat^{-1})$
where

Q $\Qmat$ is sparse.

Multivariate SPDEs

Useful when we have multiple interacting spatio-temporal processes
Consider the bi-variate sub-system of random forcings: $e_{1,t}(\svec), e_{2,t}(\svec)$ . Define the multi-variate Gaussian process (MVGP) as

(e 1, t (s), e 2, t (s)) = e (s) \sim M V G P (0, K (s, s', t, t'; θ K))

$(e_{1,t}(\svec), e_{2,t}(\svec)) = \bold{e}(\svec) \sim \mathcal{MVGP}(\bold{0},\bold{K}(\svec,\svec',t,t';\thetab_K))$

where $\bold{K}(\svec,\svec',t,t';\thetab_K)$ is a $2 \times 2$ cross-covariance matrix function (i.e.~ each element of $\bold{K}$ is a covariance function (Finley, 2007):

K (s, s', t, t') = [C o v (e i, t (s), e i, t' (s'))] i, j \in {1, 2}

$\bold{K}(\svec,\svec',t,t') = [Cov(e_{i,t}(\svec), e_{i,t'}(\svec'))]_{i,j \in \{1,2\}}$

Multivariate SPDEs

Model with a system of SPDEs. Example of a three-variate system of random forcings:

⎡ ⎣ ⎢ A 1 A 2 A 3, 2 A 2, 3 A 3 ⎤ ⎦ ⎥ ⎡ ⎣ ⎢ e 1 (s) e 2 (s) e 3 (s) ⎤ ⎦ ⎥ = ⎡ ⎣ ⎢ W 1 (s) W 2 (s) W 3 (s) ⎤ ⎦ ⎥

$\begin{bmatrix} \mathcal{A}_1 & & \\ & & \mathcal{A}_2 & \mathcal{A}_{2,3} \\ & & \mathcal{A}_{3,2} & \mathcal{A}_3 \end{bmatrix} \begin{bmatrix} e_1(\svec) \\ e_2(\svec) \\ e_3(\svec) \end{bmatrix} = \begin{bmatrix} \mathcal{W}_1(\svec) \\ \mathcal{W}_2(\svec) \\ \mathcal{W}_3(\svec) \end{bmatrix}$

Summary

SPDEs are a useful way to modelling complex phenomena.
Real power comes in using finite elements to decompose and represent them as GMRFs
From now on we restrict ourselves to the linear, Gaussian case

⎡ ⎣ ⎢ ⎢ x 1, t x 2, t ⋮ ⎤ ⎦ ⎥ ⎥ = ⎡ ⎣ ⎢ ⎢ A 1 A 2, 1 ⋮ A 1, 2 A 2, 2 ⋮ \dots \dots ⋱ ⎤ ⎦ ⎥ ⎥ ⎡ ⎣ ⎢ ⎢ x 1, t - 1 x 2, t - 1 ⋮ ⎤ ⎦ ⎥ ⎥ + J t β + ⎡ ⎣ ⎢ ⎢ e 1, t e 2, t ⋮ ⎤ ⎦ ⎥ ⎥

$\begin{bmatrix} \xvec_{1,t} \\ \xvec_{2,t} \\ \vdots \end{bmatrix} = \begin{bmatrix} \Amat_{1} & \Amat_{1,2} & \dots \\ \Amat_{2,1}& \Amat_{2,2} & \dots \\ \vdots& \vdots & \ddots \\ \end{bmatrix} \begin{bmatrix} \xvec_{1,t-1} \\ \xvec_{2,t-1} \\ \vdots \end{bmatrix} + \bold{J}_t\bold{\beta} + \begin{bmatrix} \evec_{1,t} \\ \evec_{2,t} \\ \vdots \end{bmatrix}$

Kalman filter/ RTS methods do not lend themselves to sparse methods (more of that in another talk)

Construct a massive GMRF

Collapse all variates, and their spatial and temporal components into one big GMRF
Sparsity is retained and one can make use of fill-in reduction permutations when computing the Cholesky factor (Knorr-Rue and Held, 2002)

p (x, β I | Θ) = (\prod i = 1 T p (x t | x t - 1, β, Θ)) p (x 0 | β, Θ) p (β I)

$p(\xvec,\betab_I | \Thetab) = \left(\prod_{i=1}^T{p(\xvec_{t} | \xvec_{t-1},\betab,\Thetab)}\right)p(\xvec_0 | \betab, \Thetab)p(\betab_I)$

Initial distribution (Note, not the stationary one!):

x 0 | β I, Θ \sim N (J 0 β, [Q w - A T Q w A] - 1)

$\xvec_0|\betab_I,\Thetab \sim \mathcal{N}(\bold{J}_0\betab,[\Qmat_w - \Amat^T\Qmat_w\Amat]^{-1})$

ST GMRF

$(\xvec,\betab) \sim \mathcal{N}(\bold{0},\Qmat^{-1})$ where

Q = [Q 1 Q T 2 Q 2 Q 3] Q 2 = ⎡ ⎣ ⎢ ⎢ ⎢ A T Q w J 1 - Q 0 J 0 A T Q w J 2 - Q w J 1 ⋮ ⎤ ⎦ ⎥ ⎥ ⎥,

$\small \begin{array}{cc} \Qmat = \begin{bmatrix} \Qmat_1 & \Qmat_2 \\ \Qmat_2^T & \Qmat_3 \end{bmatrix} & \Qmat_2 = \begin{bmatrix} \Amat^T\Qmat_w\Jmat_{1} - \Qmat_0\Jmat_{0} \\ \Amat^T\Qmat_w\Jmat_{2} - \Qmat_w\Jmat_1 \\ \vdots\end{bmatrix}, \end{array}$

Q 1 = ⎡ ⎣ ⎢ ⎢ ⎢ ⎢ ⎢ ⎢ ⎢ A T Q w A + Q 0 - Q w A - A T Q w A T Q w A + Q w - Q w A ⋱ - A T Q w ⋱ ⋱ ⋱ ⋱ ⎤ ⎦ ⎥ ⎥ ⎥ ⎥ ⎥ ⎥ ⎥

$\small \Qmat_1 = \begin{bmatrix} \Amat^T\Qmat_w\Amat + \Qmat_0 & -\Amat^T\Qmat_w & ~~~ \\ -\Qmat_w\Amat & \Amat^T\Qmat_w\Amat + \Qmat_w& -\Amat^T\Qmat_w & ~~~ \\ & -\Qmat_w\Amat & \ddots & \ddots \\ & \ddots & \ddots & \ddots & \end{bmatrix}$

Q 3 = J T 0 Q 0 J 0 + \sum i = 1 T (J T i Q w J i) + Q β

$\small \Qmat_3 = \Jmat_0^T\Qmat_0\Jmat_0 + \sum_{i=1}^{T} (\Jmat_i^T\Qmat_w\Jmat_i) + \Qmat_\beta$

Size matters

Conditioned on the parameters, this is the simplest model possible.
$\Qmat$ is large. Example: Antarctica, 4 ST fields, 7 time points, 500k observations, 500k vertices.
Updated $\Qmat^* \approx 3.5Gb$ . $\Qmat^{*-1} \approx 1Tb$
What if we want to find $diag(\Amat \Sigma^* \Amat^T)$ where $\Sigma^* = \Qmat^{*-1}$ ?
Impossible to carry out on a standard PC for large, $\Amat$ , even if $chol(\Qmat^*)$ is known and fill-in reduced.

Sparse subsets

A full backsolve for $\Amat \Sigma^* \Amat^T$ is wasteful since $\Amat$ is very sparse.
We only need a few elements of $\Sigma^*$ . Which elements are these?

Lemma 1: Takahashi Equations (Erisman and Tinney, 1975) Let $\Lmat = chol(\Qmat)$ . If $L_{ij} \ne 0$ , then $\Sigma_{ij}$ can be computed using elements in

$\Lmat$ and

$\Sigma^*_{IJ}, I = \{ i'; i' > i\}, J = \{j'; j' > j\}$

Knowledge of $chol(\Qmat^*)$ implies knowledge of a sparse subset of $\Sigma^*$ , $\Sigma_{00}^*$ .

Sparse subsets

A full backsolve for $\Amat \Sigma^* \Amat^T$ is wasteful since $\Amat$ is very sparse.
We only need a few elements of $\Sigma^*$ . Which elements are these?

Theorem 1: Let $\yvec = \Cmat\xvec + \vvec$ and $\vvec \sim \mathcal{N}(0,\Tmat^{-1})$ . Then if

$\textrm{zeros}(\Amat) \subseteq \textrm{zeros}(\Cmat)$

$\Amat$ is non-negative

$\Tmat$ is diagonal, then $diag(\Amat \Sigma^* \Amat^T) \equiv diag(\Amat \Sigma^*_{00} \Amat^T)$

We can replace a very expensive backsolve by a very easy matrix multiplication.

Sketch of proof

The updated precision matrix $\Qmat^* = \Qmat + \Amat^T\Tmat^{-1}\Amat$ is sparse if both $\Qmat$ and $\Amat$ are sparse and $\Tmat$ is diagonal.

Lemma 2: If $\Amat$ is non-negative then $[\Amat^T\Amat]_{jk} = 0 \Leftrightarrow diag(\Amat\Sigma^*\Amat^T)$ does not depend on $\Sigma^*_{jk}$

Let $zeros(\cdot)$ denote the symbolic pattern of a matrix.

z e r o s (Q *) = z e r o s (Q + A T T - 1 A) \geq z e r o s (A T T - 1 A) = z e r o s (A T A)

$\begin{align} zeros(\Qmat^*) &= zeros(\Qmat + \Amat^T\Tmat^{-1}\Amat) \\ &\ge zeros(\Amat^T\Tmat^{-1}\Amat) = zeros(\Amat^T\Amat) \end{align}$

Sketch of proof

By Rue and Held (2005, Corollary 2.2)

z e r o s (L *) j k \geq z e r o s (Q *) j k \geq z e r o s (A T A)

$zeros(\Lmat^*)_{jk} \ge zeros(\Qmat^*)_{jk} \\ \ge zeros(\Amat^T\Amat)$

Proof is completed by applying Lemma 1:

z e r o s (Σ 00) \geq z e r o s (L *) j k

$zeros(\Sigma_{00}) \ge zeros(\Lmat^*)_{jk}$

Implications

In practice the Takahashi equations are always used with GMRFs in order to obtain the marginal variances. Theorem 2 provides a means to obtain uncertainty on linear combinations essentially for free.
I can provide predictive uncertainties at any point provided that each row of $zeros(\Amat) \subseteq zeros(\Cmat)$ .

Implications

Hundreds of thousands of validation points

(Pure) parallel sampling with massive GMRFs

In all practical cases there are unknown parameters and a massive GMRF, sampling may be required
MCMC with multiple chains is relatively straightforward
- e.g. Evolutionary MCMC, Embarassingly parallel MCMC
Extremely little work on parallelising a single chain
Take advantage of local Markov property

x i ⊥ x - {i, n e (i)} | x n e (i) \forall i \in V

$x_i \perp \xvec_{-\{i,ne(i)\}} | \xvec_{ne(i)} ~~~ \forall i \in \mathcal{V}$

E (x i | x - i) = - 1 Q i i \sum j : j \sim i Q i j x j; p r e c (x i | x - i) = Q i i

$E(x_i | \xvec_{-i}) = -\frac{1}{Q_{ii}}\sum_{j:j\sim i}Q_{ij}x_j; ~~~ prec(x_i | \xvec_{-i}) = Q_{ii}$

Graph colouring

Theorem 3: Any planar map divided into contigious regions can be coloured using at most FOUR colours (Appel and Haken, 1976)

Graph colouring

Theorem 3: Any planar map divided into contigious regions can be coloured using at most FOUR colours (Appel and Haken, 1976)

Parallel (chromatic) Sampling

Improving the chromatic sampler

Single-site updating in spatio-temporal systems is a bad idea
We can make use of the global Markov property

x A ⊥ x B | x C

$\xvec_A \perp \xvec_B | \xvec_{C}$

$\forall A,B,C \notin \varnothing$ and $C$ separating $A,B$ .

In typical Bayesian networks, finding $A, B$ and $C$ can be difficult (seed and spawn idea, Gonzales, 2011).
With a spatial network, this is remarkably easy using the Kernighan-Lin algorithm for graph partitioning (Kernighan and Lin, 1970).

Kernighan-Lin algorithm

Let $G(V,E)$ be a graph, and let $V$ be the set of nodes and $E$ the set of edges.
Find a partition of two disjoint, equal-size subsets $A,B$ , such that the sum of the weights of the edges between nodes in A and B is minimized.
Setting edge weights = 1, gives a bisection algorithm!

Satellite footprints

Sampling from the prior guarantees enormous speed-ups using MPI
Observations with large spatial footprints cause the four-colour theorem to fail.

Q ˆ i i x^i = C i T Q ⎛ ⎝ y - \sum j \in [1 \dots d] / i C j x^j ⎞ ⎠ - \sum j \in [1 \dots d] / i (Q i j x^j)

$\begin{align} \widehat\Qmat^{ii} \hat\xvec^i &= \Cmat^{i^T}\Qmat\left(\yvec - \sum_{j \in [1\dots d]/i}\Cmat^j\hat\xvec^j\right) - \\ &~~~~ \sum_{j \in [1\dots d]/i} \left(\Qmat^{ij}\hat\xvec^j \right) \end{align}$

Sparse linear algebraic methods and MPI for spatio-temporal inference

Introduction

Introduction

Introduction

Bayesian Hierarchical Model

Bayesian Hierarchical Model

Modelling approaches

Dynamic process models

Dynamic process models

Moving on to discrete time

SPDE link to Matern fields

Why SPDEs

Dimensionality reduction

Dimensionality reduction

Modelling on manifolds

Taking into consideration boundary effects

Multivariate SPDEs

Multivariate SPDEs

Summary

Construct a massive GMRF

ST GMRF

Size matters

The medal plot

The medal plot

Sparse subsets

Sparse subsets

Sketch of proof

Sketch of proof

Demo

Implications

Implications

(Pure) parallel sampling with massive GMRFs

Graph colouring

Graph colouring

Parallel (chromatic) Sampling

Improving the chromatic sampler

Kernighan-Lin algorithm

Example

Example

Example

Example

Example

Example

Example

Example

Satellite footprints

New colouring strategy

New colouring strategy

Conclusion

Thanks